import torch
import datetime

from utils import *
from sparseTools.utils import SparseAdjList


def adjsProcess(adjs):
    # 遍历邻接矩阵adjs的每个键值对k
    # 这句代码的整体作用是将邻接矩阵中每条边的权重重新设置为该边起始节点度的倒数。这样做的目的是进行归一化处理，使得每个节点连接的边权重之和为1，通常用于消息传递中的归一化步骤。
    # todo 这里不知道对于非 Freebase 数据集是否有作用，待检查（对于非 Freebase 数据集，使用列表存储；对于 Freebase 数据集，使用字典存储。）
    for k in adjs.keys():
        # 清除当前邻接矩阵adjs[k]的权重值。
        adjs[k].storage._value = None
        # 设置新的权重值，所有边的权重设置为平均值。
        # adjs[k] 是一个稀疏矩阵（通常是 DGLGraph 的邻接矩阵）。
        # nnz() 函数返回稀疏矩阵中非零元素的数量（number of non-zero elements）。
        # torch.ones 函数创建一个包含全为1的张量，其长度为 adjs[k].nnz()，即包含和 adjs[k] 中非零元素数量相同的1。
        # 计算邻接矩阵 adjs[k] 每行的和，结果是一个一维张量，其中每个元素代表邻接矩阵对应行的和。
        # dim=-1 表示在最后一个维度上进行求和，这里指的是行求和。
        # adjs[k].storage 是稀疏矩阵的底层存储结构，row() 返回一个张量，表示稀疏矩阵中每个非零元素的行索引。
        # 通过使用行索引张量 adjs[k].storage.row() 访问按行求和的结果，这会为每个非零元素选择其对应的行和。
        # 结果是一个张量，其长度等于 adjs[k] 中非零元素的数量，每个元素是相应非零元素所在行的和。
        # 对于每个非零元素，计算1除以该元素所在行的和。
        # 结果是一个张量，其长度为 adjs[k].nnz()，每个元素代表邻接矩阵中每条边的权重。
        # 最终，将计算得到的权重张量分配给 adjs[k].storage._value，这将更新邻接矩阵中每条边的权重。
        adjs[k].storage._value = torch.ones(adjs[k].nnz()) / adjs[k].sum(dim=-1)[adjs[k].storage.row()]

    return adjs


def neighboorAggregation(args, g, dl, extra_metapath, tgt_type, adjs, num_tgt_nodes, prop_device):
    # 根据数据集不同，计算k跳特征。
    # 如果数据集不是Freebase，则使用hg_propagate_feat_dgl函数传播特征，并提取节点特征。
    # 如果数据集是Freebase，则根据是否已有保存的特征邻接矩阵加载或生成特征邻接矩阵，并提取节点特征。
    # 获取当前时间，标记特征传播开始的时间点。
    prop_tic = datetime.datetime.now()
    if args.dataset != 'Freebase':
        # 如果有额外的元路径 (extra_metapath)，则计算这些元路径的最大长度 max_length，以 args.num_hops + 1 和 extra_metapath 中最大长度的较大值为准。
        # 如果没有 extra_metapath，则直接使用 args.num_hops + 1 作为 max_length。
        if len(extra_metapath):
            max_length = max(args.num_hops + 1, max([len(ele) for ele in extra_metapath]))
        else:
            max_length = args.num_hops + 1

        # 使用 hg_propagate_feat_dgl 函数对图 g 进行特征传播，传播的跳数为 args.num_hops，并考虑 max_length 和 extra_metapath。
        g = hg_propagate_feat_dgl(g, tgt_type, args.num_hops, max_length, extra_metapath, echo=True)
        # 初始化一个空字典 raw_feats 用于存储处理后的特征。
        raw_feats = {}
        # 获取图 g 中目标类型节点的所有特征键。
        keys = list(g.nodes[tgt_type].data.keys())
        # 将图 g 中目标类型节点的数据（特征）弹出，并存储到 raw_feats 中。
        for k in keys:
            raw_feats[k] = g.nodes[tgt_type].data.pop(k)
    else:
        # 和非 Freebase 数据集相同，计算 max_length。
        if len(extra_metapath):
            max_length = max(args.num_hops + 1, max([len(ele) for ele in extra_metapath]))
        else:
            max_length = args.num_hops + 1

        # 如果跳数为1，则从邻接矩阵 adjs 中克隆出目标类型的邻接矩阵，并存储在 raw_meta_adjs 字典中。
        if args.num_hops == 1:
            raw_meta_adjs = {k: v.clone() for k, v in adjs.items() if k[0] == tgt_type}
        else:
            # 如果跳数不为1，则准备保存特征邻接矩阵的文件名 save_name，并检查是否已经存在该文件。
            save_name = f'./Freebase_adjs/feat_hop{args.num_hops}'
            # 如果文件存在，则加载保存的特征邻接矩阵。通过 SparseAdjList 类加载，并更新到 raw_meta_adjs 字典中。
            if os.path.exists(f'{save_name}_00_int64.npy'):
                raw_meta_adjs = {}
                for srcname in dl.nodes['count'].keys():
                    print(f'Loading feature adjs from {save_name}_0{srcname}')
                    tmp = SparseAdjList(f'{save_name}_0{srcname}', None, None, num_tgt_nodes,
                                        dl.nodes['count'][srcname], with_values=True)
                    for k in tmp.keys:
                        assert k not in raw_meta_adjs
                    raw_meta_adjs.update(tmp.load_adjs(expand=True))
                    del tmp
            else:
                # 如果文件不存在，则生成特征邻接矩阵，调用 hg_propagate_sparse_pyg 函数进行特征传播。
                print(
                    'Generating feature adjs for Freebase ...\n(For each configutaion, this happens only once for the first execution, and results will be saved for later executions)')
                raw_meta_adjs = hg_propagate_sparse_pyg(adjs, tgt_type, args.num_hops, max_length, extra_metapath,
                                                        prop_feats=True, echo=True, prop_device=prop_device)

                # 将生成的特征邻接矩阵保存到文件中，以便后续使用。
                meta_adj_list = []
                for srcname in dl.nodes['count'].keys():
                    keys = [k for k in raw_meta_adjs.keys() if k[-1] == str(srcname)]
                    print(f'Saving feature adjs {keys} into {save_name}_0{srcname}')
                    tmp = SparseAdjList(f'{save_name}_0{srcname}', keys, raw_meta_adjs, num_tgt_nodes,
                                        dl.nodes['count'][srcname], with_values=True)
                    meta_adj_list.append(tmp)

                # 加载保存的特征邻接矩阵，并进行一致性检查，确保加载的数据和生成的数据一致。
                for srcname in dl.nodes['count'].keys():
                    tmp = SparseAdjList(f'{save_name}_0{srcname}', None, None, num_tgt_nodes,
                                        dl.nodes['count'][srcname], with_values=True)
                    tmp_adjs = tmp.load_adjs(expand=True)
                    print(srcname, tmp.keys)
                    for k in tmp.keys:
                        assert torch.all(raw_meta_adjs[k].storage.rowptr() == tmp_adjs[k].storage.rowptr())
                        assert torch.all(raw_meta_adjs[k].storage.col() == tmp_adjs[k].storage.col())
                        assert torch.all(raw_meta_adjs[k].storage.value() == tmp_adjs[k].storage.value())
                    del tmp_adjs, tmp
                    gc.collect()

        # 复制 raw_meta_adjs 中的邻接矩阵特征到 raw_feats 中，只包括跳数小于等于 args.num_hops + 1 或在 extra_metapath 中的邻接矩阵。
        raw_feats = {k: v.clone() for k, v in raw_meta_adjs.items() if
                     len(k) <= args.num_hops + 1 or k in extra_metapath}

        # 确保 raw_feats 中没有键 '0'，然后添加一个单位矩阵（对角矩阵），表示自连接。
        assert '0' not in raw_feats
        raw_feats['0'] = SparseTensor.eye(dl.nodes['count'][0])

    # 打印目标节点类型tgt_type的特征键值及其数量。
    # 对于Freebase数据集，以稀疏张量模式打印特征键值及其尺寸。
    # 对于其他数据集，打印特征键值及其尺寸。
    print(f'For tgt type {tgt_type}, feature keys (num={len(raw_feats)}):', end='')
    if args.dataset == 'Freebase':
        print(' (in SparseTensor mode)')
        for k, v in raw_feats.items():
            print(k, v.sizes())
    else:
        print()
        for k, v in raw_feats.items():
            print(k, v.size())
    print()

    # 根据数据集类型，设置数据尺寸data_size。
    # prop_toc = datetime.datetime.now(): 记录特征传播结束时间。
    # 打印预处理所花费的时间。
    if args.dataset in ['DBLP', 'ACM', 'IMDB']:
        data_size = {k: v.size(-1) for k, v in raw_feats.items()}
    elif args.dataset == 'Freebase':
        data_size = dict(dl.nodes['count'])
    else:
        assert 0
    prop_toc = datetime.datetime.now()
    print(f'Time used for feat prop {prop_toc - prop_tic}')
    gc.collect()

    return raw_feats, data_size